//===--- LSLocationPrinter.cpp - Dump all memory locations in program -----===//
//
// This source file is part of the Swift.org open source project
//
// Copyright (c) 2014 - 2017 Apple Inc. and the Swift project authors
// Licensed under Apache License v2.0 with Runtime Library Exception
//
// See https://swift.org/LICENSE.txt for license information
// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
//
//===----------------------------------------------------------------------===//
//
// This pass tests type expansion, memlocation expansion and memlocation
// reduction.
//
//===----------------------------------------------------------------------===//

#define DEBUG_TYPE "pil-memlocation-dumper"

#include "polarphp/pil/lang/Projection.h"
#include "polarphp/pil/lang/PILFunction.h"
#include "polarphp/pil/lang/PILValue.h"
#include "polarphp/pil/optimizer/passmgr/Passes.h"
#include "polarphp/pil/optimizer/passmgr/Transforms.h"
#include "polarphp/pil/optimizer/utils/LoadStoreOptUtils.h"
#include "llvm/Support/CommandLine.h"

using namespace polar;

//===----------------------------------------------------------------------===//
//                              Top Level Driver
//===----------------------------------------------------------------------===//

namespace {

enum class MLKind : unsigned {
   OnlyExpansion = 0,
   OnlyReduction = 1,
   OnlyTypeExpansion = 2,
   All = 3,
};

} // end anonymous namespace

static llvm::cl::opt<MLKind> LSLocationKinds(
   "ml", llvm::cl::desc("LSLocation Kinds:"), llvm::cl::init(MLKind::All),
   llvm::cl::values(
      clEnumValN(MLKind::OnlyExpansion, "only-expansion", "only-expansion"),
      clEnumValN(MLKind::OnlyReduction, "only-reduction", "only-reduction"),
      clEnumValN(MLKind::OnlyTypeExpansion, "only-type-expansion",
                 "only-type-expansion"),
      clEnumValN(MLKind::All, "all", "all")));

static llvm::cl::opt<bool> UseProjection("lslocation-dump-use-new-projection",
                                         llvm::cl::init(false));

namespace {

class LSLocationPrinter : public PILModuleTransform {
   /// Type expansion analysis.
   TypeExpansionAnalysis *TE;

public:
   /// Dumps the expansions of PILType accessed in the function.
   /// This tests the expandTypeIntoLeafProjectionPaths function, which is
   /// a function used extensively in expand and reduce functions.
   ///
   /// We test it to catch any suspicious things in the earliest point.
   ///
   void printTypeExpansion(PILFunction &Fn) {
      PILModule *M = &Fn.getModule();
      ProjectionPathList PPList;
      unsigned Counter = 0;
      for (auto &BB : Fn) {
         for (auto &II : BB) {
            if (auto *LI = dyn_cast<LoadInst>(&II)) {
               PILValue V = LI->getOperand();
               // This is an address type, take it object type.
               PILType Ty = V->getType().getObjectType();
               ProjectionPath::expandTypeIntoLeafProjectionPaths(
                  Ty, M, TypeExpansionContext(Fn), PPList);
            } else if (auto *SI = dyn_cast<StoreInst>(&II)) {
               PILValue V = SI->getDest();
               // This is an address type, take it object type.
               PILType Ty = V->getType().getObjectType();
               ProjectionPath::expandTypeIntoLeafProjectionPaths(
                  Ty, M, TypeExpansionContext(Fn), PPList);
            } else {
               // Not interested in these instructions yet.
               continue;
            }

            llvm::outs() << "#" << Counter++ << II;
            for (auto &T : PPList) {
               T.getValue().print(llvm::outs(), *M, TypeExpansionContext(Fn));
            }
            PPList.clear();
         }
      }
      llvm::outs() << "\n";
   }

   void printTypeExpansionWithProjection(PILFunction &Fn) {
      PILModule *M = &Fn.getModule();
      llvm::SmallVector<Optional<ProjectionPath>, 8> PPList;
      unsigned Counter = 0;
      for (auto &BB : Fn) {
         for (auto &II : BB) {
            PILValue V;
            PILType Ty;
            if (auto *LI = dyn_cast<LoadInst>(&II)) {
               V = LI->getOperand();
               // This is an address type, take it object type.
               Ty = V->getType().getObjectType();
               ProjectionPath::expandTypeIntoLeafProjectionPaths(
                  Ty, M, TypeExpansionContext(Fn), PPList);
            } else if (auto *SI = dyn_cast<StoreInst>(&II)) {
               V = SI->getDest();
               // This is an address type, take it object type.
               Ty = V->getType().getObjectType();
               ProjectionPath::expandTypeIntoLeafProjectionPaths(
                  Ty, M, TypeExpansionContext(Fn), PPList);
            } else {
               // Not interested in these instructions yet.
               continue;
            }

            llvm::outs() << "#" << Counter++ << II;
            for (auto &T : PPList) {
               T.getValue().print(llvm::outs(), *M, TypeExpansionContext(Fn));
            }
            PPList.clear();
         }
      }
      llvm::outs() << "\n";
   }

   /// Dumps the expansions of memory locations accessed in the function.
   /// This tests the expand function in LSLocation class.
   ///
   /// We test it to catch any suspicious things when memory location is
   /// expanded, i.e. base is traced back and aggregate is expanded
   /// properly.
   void printMemExpansion(PILFunction &Fn) {
      LSLocation L;
      LSLocationList Locs;
      unsigned Counter = 0;
      for (auto &BB : Fn) {
         for (auto &II : BB) {
            if (auto *LI = dyn_cast<LoadInst>(&II)) {
               PILValue Mem = LI->getOperand();
               PILValue UO = getUnderlyingObject(Mem);
               L.init(UO, ProjectionPath::getProjectionPath(UO, Mem));
               if (!L.isValid())
                  continue;
               LSLocation::expand(L, &Fn.getModule(), TypeExpansionContext(Fn), Locs,
                                  TE);
            } else if (auto *SI = dyn_cast<StoreInst>(&II)) {
               PILValue Mem = SI->getDest();
               PILValue UO = getUnderlyingObject(Mem);
               L.init(UO, ProjectionPath::getProjectionPath(UO, Mem));
               if (!L.isValid())
                  continue;
               LSLocation::expand(L, &Fn.getModule(), TypeExpansionContext(Fn), Locs,
                                  TE);
            } else {
               // Not interested in these instructions yet.
               continue;
            }

            llvm::outs() << "#" << Counter++ << II;
            for (auto &Loc : Locs) {
               Loc.print(llvm::outs(), &Fn.getModule(), TypeExpansionContext(Fn));
            }
            Locs.clear();
         }
      }
      llvm::outs() << "\n";
   }

   /// Dumps the reductions of set of memory locations.
   ///
   /// This function first calls expand on a memory location. It then calls
   /// reduce, in hope to get the original memory location back.
   ///
   void printMemReduction(PILFunction &Fn) {
      LSLocation L;
      LSLocationList Locs;
      LSLocationList SLocs;
      unsigned Counter = 0;
      for (auto &BB : Fn) {
         for (auto &II : BB) {

            // Expand it first.
            //
            if (auto *LI = dyn_cast<LoadInst>(&II)) {
               PILValue Mem = LI->getOperand();
               PILValue UO = getUnderlyingObject(Mem);
               L.init(UO, ProjectionPath::getProjectionPath(UO, Mem));
               if (!L.isValid())
                  continue;
               LSLocation::expand(L, &Fn.getModule(), TypeExpansionContext(Fn), Locs,
                                  TE);
            } else if (auto *SI = dyn_cast<StoreInst>(&II)) {
               PILValue Mem = SI->getDest();
               PILValue UO = getUnderlyingObject(Mem);
               L.init(UO, ProjectionPath::getProjectionPath(UO, Mem));
               if (!L.isValid())
                  continue;
               LSLocation::expand(L, &Fn.getModule(), TypeExpansionContext(Fn), Locs,
                                  TE);
            } else {
               // Not interested in these instructions yet.
               continue;
            }

            // Try to reduce it.
            //
            // Reduction should not care about the order of the memory locations in
            // the set.
            for (auto I = Locs.begin(); I != Locs.end(); ++I) {
               SLocs.push_back(*I);
            }

            // This should get the original (unexpanded) location back.
            LSLocation::reduce(L, &Fn.getModule(), TypeExpansionContext(Fn), SLocs);
            llvm::outs() << "#" << Counter++ << II;
            for (auto &Loc : SLocs) {
               Loc.print(llvm::outs(), &Fn.getModule(), TypeExpansionContext(Fn));
            }
            L.reset();
            Locs.clear();
            SLocs.clear();
         }
      }
      llvm::outs() << "\n";
   }

   void run() override {
      for (auto &Fn : *getModule()) {
         if (Fn.isExternalDeclaration()) continue;

         // Initialize the type expansion analysis.
         TE = PM->getAnalysis<TypeExpansionAnalysis>();

         llvm::outs() << "@" << Fn.getName() << "\n";
         switch (LSLocationKinds) {
            case MLKind::OnlyTypeExpansion:
               printTypeExpansionWithProjection(Fn);
               break;
            case MLKind::OnlyExpansion:
               printMemExpansion(Fn);
               break;
            case MLKind::OnlyReduction:
               printMemReduction(Fn);
               break;
            default:
               break;
         }
      }
   }

};

} // end anonymous namespace

PILTransform *polar::createLSLocationPrinter() {
   return new LSLocationPrinter();
}
